from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage
import json

class BasicState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]

class BasicGraphDemo:
    """基础图执行演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建基础图"""
        
        # 定义工作流
        workflow = StateGraph(BasicState)
        
        # 添加节点
        workflow.add_node("process_input", self._process_input)
        workflow.add_node("generate_response", self._generate_response)
        workflow.add_node("format_output", self._format_output)
        
        # 设置入口点
        workflow.set_entry_point("process_input")
        
        # 设置边
        workflow.add_edge("process_input", "generate_response")
        workflow.add_edge("generate_response", "format_output")
        workflow.add_edge("format_output", END)
        
        return workflow.compile()
    
    def _process_input(self, state: BasicState) -> BasicState:
        """处理输入"""
        messages = state["messages"]
        if not messages:
            messages = [HumanMessage(content="你好，请介绍一下你自己")]
        
        # 模拟处理逻辑
        processed_content = f"收到消息: {messages[-1].content}"
        
        return {
            "messages": messages,
            "result": processed_content
        }
    
    def _generate_response(self, state: BasicState) -> BasicState:
        """生成响应"""
        processed_content = state["result"]
        
        # 模拟AI响应生成
        response = f"AI助手回复: 我理解您说的是 '{processed_content}'，我会为您提供帮助。"
        
        # 添加AI消息到对话历史
        messages = state["messages"] + [AIMessage(content=response)]
        
        return {
            "messages": messages,
            "result": response
        }
    
    def _format_output(self, state: BasicState) -> BasicState:
        """格式化输出"""
        result = state["result"]
        
        # 格式化最终结果
        formatted_result = f"最终结果: {result}"
        
        return {
            "messages": state["messages"],
            "result": formatted_result
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行基础图"""
        # 准备初始状态
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "你好"))],
            "result": ""
        }
        
        # 执行图
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "messages": [msg.content for msg in result["messages"]],
            "final_result": result["result"]
        }
